from flask_socketio import SocketIO
from collections import deque
from dataclasses import dataclass
from typing import Optional

from .nodes.agentGraphNode import AgantGraphNode
from .nodes.baseGraphNode import BaseGraphNode
from .nodes.httpGraphNode import HttpGraphNode
from .nodes.sessionInfo import SessionInfo
from .nodes.toolExecuteNode import ToolExecuteNode
from .nodes.toolNode import ToolNode
from .nodes.httpInGraphNode import HttpInGraphNode
from .nodes.httpOutGraphNode import HttpOutGraphNode
from .nodes.httpAudioInGraphNode import HttpAudioInGraphNode
from .nodes.sttGraphNode import STTGraphNode
@dataclass
class ExecutionTask:
    node_id: str
    session_info: SessionInfo

class GraphNodeManager:
    def __init__(self, project):
        from .project import Project
        assert isinstance(project, Project), "project must be an instance of Project"
        self.nodesByIdDict = {}
        self.edges = None
        self.nodes = None
        self.project = project
        self.sid_to_sessionInfo = {}
        # 执行队列
        self.execution_queue = deque()
        # 记录队列中的节点ID，用于防止重复入队
        self.nodes_in_queue = set()
        # 是否正在处理队列
        self.is_processing_queue = False
        self.http_node = None

    def stop(self):
        for k,v in self.nodesByIdDict.items():
            v.stop()
        # 清空执行队列
        self.execution_queue.clear()
        self.nodes_in_queue.clear()
        self.is_processing_queue = False

    def add_to_execution_queue(self, node_id: str, session_info: SessionInfo) -> None:
        """添加执行任务到队列"""
        if node_id not in self.nodes_in_queue:
            self.execution_queue.append(ExecutionTask(node_id, session_info))
            self.nodes_in_queue.add(node_id)
            # 如果队列没有在处理，则开始处理
            if not self.is_processing_queue:
                self.process_execution_queue()

    def process_execution_queue(self) -> None:
        """处理执行队列"""
        self.is_processing_queue = True
        while self.execution_queue:
            task = self.execution_queue.popleft()
            self.nodes_in_queue.remove(task.node_id)
            try:
                # 执行节点
                self.nodesByIdDict[task.node_id].execute(task.session_info)
            except Exception as e:
                print(f"Error executing node {task.node_id}: {str(e)}")
        self.is_processing_queue = False

    def deploy(self, nodes, edges, random_key:str):
        self.nodes = nodes
        self.edges = edges
        for node in nodes:
            if node['type'] == 'http': # 先创建http节点，因为http节点需要创建ws服务
                self.nodesByIdDict[node['id']] = HttpGraphNode(node['id'], node, self, random_key)
                self.http_node = self.nodesByIdDict[node['id']]
        for node in nodes:
            if node['type'] == 'http':
                continue
            elif node['type'] == 'agent':
                self.nodesByIdDict[node['id']] = AgantGraphNode(node['id'], node, self, random_key)
            elif node['type']=='tool-execute':
                self.nodesByIdDict[node['id']] = ToolExecuteNode(node['id'], node, self, random_key)
            elif node['type']=='tool':
                self.nodesByIdDict[node['id']] = ToolNode(node['id'], node, self, random_key)
            elif node['type']=='http-in':
                self.nodesByIdDict[node['id']] = HttpInGraphNode(node['id'], node, self, random_key)
            elif node['type']=='http-out':
                self.nodesByIdDict[node['id']] = HttpOutGraphNode(node['id'], node, self, random_key)   
            elif node['type']=='http-audio-in':
                self.nodesByIdDict[node['id']] = HttpAudioInGraphNode(node['id'], node, self, random_key)
            elif node['type']=='stt':
                self.nodesByIdDict[node['id']] = STTGraphNode(node['id'], node, self, random_key)
            else:
                self.nodesByIdDict[node['id']] = BaseGraphNode(node['id'], node, self, random_key)

        # 依据flow边建立Topology规则
        for edge in edges:
            sourceHandle = edge['sourceHandle']
            targetHandle = edge['targetHandle']
            sourceID = edge['source']
            targetID = edge['target']
            if ('flow' in sourceHandle or 'flow' in targetHandle):
                self.nodesByIdDict[sourceID].add_flow_next(targetID, sourceHandle=sourceHandle)
            # 依据数据边建立数据取用规则
            else:
                self.nodesByIdDict[targetID].dict['in'][targetHandle] = {'nodeId': sourceID, 'handle': sourceHandle}

        # 基于HTTP节点属性建立WS服务
        for node in nodes:
            if node['type'] == 'http':
                self.add_to_execution_queue(node['id'], None)
            elif node['type'] == 'http-in':
                self.add_to_execution_queue(node['id'], None)
            elif node['type'] == 'http-audio-in':
                self.add_to_execution_queue(node['id'], None)

    def createSessionInfo(self, session_id:str, ws: SocketIO, sid: str, route:str, full_url:str):
        if(self.nodes is None) or (self.edges is None): return None
        sessionInfo = SessionInfo(session_id,ws, sid, route, full_url)
        self.sid_to_sessionInfo[sid] = sessionInfo
        print(f'SessionInfo created for sid: {sid}')
        for nk,nv in self.nodesByIdDict.items():
            print(f'node:{nk}, flow_next_nodes:{nv.flow_next_nodes}')
            for next_node_id in nv.get_all_flow_next_nodes():
                if next_node_id not in sessionInfo.flow_prev_nodes_executed_number:
                    sessionInfo.flow_prev_nodes_executed_number[next_node_id] = {}
                sessionInfo.flow_prev_nodes_executed_number[next_node_id][nk] = 0
        return sessionInfo
